#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelAddcmul
{
public:
  __aicore__ inline KernelAddcmul() {}
  __aicore__ inline void Init(GM_ADDR input_data, GM_ADDR x1, GM_ADDR x2, GM_ADDR value, GM_ADDR y, uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t typeKey, uint32_t x2Length)
  {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    this->typeKey = typeKey;
    this->value = *((__gm__ DTYPE_X1 *)value);
    this->blockLength = blockLength;
    this->lasttileLength = lasttileLength;
    this->x2Length = x2Length;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    input_dataGm.SetGlobalBuffer((__gm__ DTYPE_X1 *)input_data,
                                this->blockLength);
    x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x1,
                        this->blockLength);
    x2Gm.SetGlobalBuffer((__gm__ DTYPE_X1 *)x2,
                        this->blockLength);
    yGm.SetGlobalBuffer((__gm__ DTYPE_X1 *)y,
                          this->blockLength);
    
    if(this->x2Length == 0){
      this->tileLength = tileLength / BUFFER_NUM;
      pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * 3 * sizeof(DTYPE_X1));
      pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X1));
    }
    else{
      this->tileLength = tileLength;
      pipe.InitBuffer(inQueueIN_1, 1, this->tileLength * 3 * sizeof(DTYPE_X1));
      pipe.InitBuffer(outQueueY_1, 1, this->tileLength * sizeof(DTYPE_X1));
    }
    if(this->typeKey == 2){
      pipe.InitBuffer(calcBuf, this->tileLength * 2 * sizeof(half));
      pipe.InitBuffer(calcBuf2, this->tileLength * sizeof(int32_t));
    }
  }
  __aicore__ inline void Process()
  {
    if(this->x2Length == 0){
      int32_t loopCount = this->tileNum * BUFFER_NUM;
      for (int32_t i = 0; i < loopCount; i++)
      {
        if(BUFFER_NUM == 2 && i == loopCount - 2 && lasttileLength <= tileLength){
          continue;
        }
        CopyIn(i);
        Compute(i);
        CopyOut(i);
      }
    }
    else {
      LocalTensor<DTYPE_X1> inLocal = inQueueIN_1.AllocTensor<DTYPE_X1>();
      DataCopy(inLocal[0], input_dataGm[0], (this->tileLength));
      DataCopy(inLocal[this->tileLength], x1Gm[0], (this->tileLength));
      DataCopy(inLocal[this->tileLength * 2], x2Gm[0], this->x2Length);
      inQueueIN_1.EnQue(inLocal);
      
      Compute(0);

      LocalTensor<DTYPE_X1> yLocal = outQueueY_1.DeQue<DTYPE_X1>();
      DataCopy(yGm[0], yLocal, (this->tileLength));
      outQueueY_1.FreeTensor(yLocal);
    }
  }

private:
  __aicore__ inline void CopyIn(int32_t progress)
  {
    LocalTensor<DTYPE_X1> inLocal = inQueueIN.AllocTensor<DTYPE_X1>();
    if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
            inLocal[0],
            input_dataGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
        DataCopy(
            inLocal[this->tileLength],
            x1Gm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
        DataCopy(
            inLocal[2 * (this->tileLength)],
            x2Gm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));

      }
      else {
        DataCopy(inLocal[0], input_dataGm[progress * (this->tileLength)],
                 (this->tileLength));
        DataCopy(inLocal[this->tileLength], x1Gm[progress * this->tileLength],
                 this->tileLength);
        DataCopy(inLocal[2 * (this->tileLength)],
                 x2Gm[progress * this->tileLength], this->tileLength);
      }
    inQueueIN.EnQue(inLocal);
    
  }
  __aicore__ inline void Compute(int32_t progress)
  {
    if(this->typeKey == 2){
      LocalTensor<int8_t> yLocal, inLocal;
      if(x2Length != 0){
        inLocal = inQueueIN_1.DeQue<int8_t>();
        yLocal = outQueueY_1.AllocTensor<int8_t>();
      }else{
        inLocal = inQueueIN.DeQue<int8_t>();
        yLocal = outQueueY.AllocTensor<int8_t>();
      }
      LocalTensor<int8_t> input_dataLocal = inLocal[0];
      LocalTensor<int8_t> x1Local = inLocal[this->tileLength];
      LocalTensor<int8_t> x2Local = inLocal[2 * (this->tileLength)];
      LocalTensor<half> tmpLocal = calcBuf.Get<half>();
      LocalTensor<half> x1Local_half = tmpLocal;
      LocalTensor<half> x2Local_half = tmpLocal[this->tileLength];
      LocalTensor<int32_t> turnLocal = calcBuf2.Get<int32_t>();

      Cast(x1Local_half, x1Local, RoundMode::CAST_NONE, this->tileLength);
      Cast(x2Local_half, x2Local, RoundMode::CAST_NONE, this->tileLength);
      if(x2Length != 0){
        for(int i = this->x2Length; i<this->tileLength; i++){
          x2Local_half.SetValue(i, x2Local_half.GetValue(i % this->x2Length));
        }
      }
      Mul(x1Local_half, x1Local_half, x2Local_half, this->tileLength);

      Muls(x1Local_half, x1Local_half, (half)this->value, this->tileLength);
      
      Cast(x2Local_half, input_dataLocal, RoundMode::CAST_NONE, this->tileLength);

      Add(x1Local_half, x1Local_half, x2Local_half, this->tileLength);
      Cast(turnLocal, x1Local_half, RoundMode::CAST_RINT, this->tileLength);
      for(int32_t i = 0; i<this->tileLength; i++){
        turnLocal.SetValue(i, turnLocal.GetValue(i) % (int32_t)256);
        if(turnLocal.GetValue(i) > (int32_t)127){
          turnLocal.SetValue(i, turnLocal.GetValue(i) - (int32_t)256);
        }
        if(turnLocal.GetValue(i) < (int32_t)-128){
          turnLocal.SetValue(i, (int32_t)256 + turnLocal.GetValue(i));
        }
      }
      Cast(x1Local_half, turnLocal, RoundMode::CAST_RINT, this->tileLength);
      Cast(yLocal, x1Local_half, RoundMode::CAST_NONE, this->tileLength);
      if(x2Length != 0){
        outQueueY_1.EnQue<int8_t>(yLocal);
        inQueueIN_1.FreeTensor(inLocal);
      }else{
        outQueueY.EnQue<int8_t>(yLocal);
        inQueueIN.FreeTensor(inLocal);
      }
      
    }else{
      LocalTensor<DTYPE_X1> yLocal, inLocal;
      if(x2Length != 0){
        inLocal = inQueueIN_1.DeQue<DTYPE_X1>();
        yLocal = outQueueY_1.AllocTensor<DTYPE_X1>();
      }else{
        inLocal = inQueueIN.DeQue<DTYPE_X1>();
        yLocal = outQueueY.AllocTensor<DTYPE_X1>();
      }
      LocalTensor<DTYPE_X1> input_dataLocal = inLocal;
      LocalTensor<DTYPE_X1> x1Local = inLocal[this->tileLength];
      LocalTensor<DTYPE_X1> x2Local = inLocal[2 * (this->tileLength)];
      // compute
      if(x2Length != 0){
        for(int32_t i = this->x2Length; i < this->tileLength; i++){
          x2Local.SetValue(i, x2Local.GetValue(i % this->x2Length));
        }
      }
      Mul(yLocal, x1Local, x2Local, this->tileLength);
      Muls(yLocal, yLocal, this->value, this->tileLength);
      Add(yLocal, input_dataLocal, yLocal, this->tileLength);
      if(x2Length != 0){
        outQueueY_1.EnQue<DTYPE_X1>(yLocal);
        inQueueIN_1.FreeTensor(inLocal);
      }else{
        outQueueY.EnQue<DTYPE_X1>(yLocal);
        inQueueIN.FreeTensor(inLocal);
      }
    }

  }
  __aicore__ inline void CopyOut(int32_t progress)
  {
    LocalTensor<DTYPE_X1> yLocal = outQueueY.DeQue<DTYPE_X1>();
    if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
        (progress == (this->tileNum * BUFFER_NUM - 1))) {
        DataCopy(
          yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
          yLocal, (this->tileLength));
    }

    else {
      DataCopy(yGm[progress * (this->tileLength)], yLocal,
                (this->tileLength));
    }
    outQueueY.FreeTensor(yLocal);
  }

private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
  TQue<QuePosition::VECIN, 1> inQueueIN_1, inQueueX1, inQueueX2;
  TBuf<TPosition::VECCALC>calcBuf, calcBuf2;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
  TQue<QuePosition::VECOUT, 1> outQueueY_1;
  GlobalTensor<DTYPE_INPUT_DATA> input_dataGm;
  GlobalTensor<DTYPE_X1> x1Gm;
  GlobalTensor<DTYPE_X2> x2Gm;
  GlobalTensor<DTYPE_Y> yGm;
  DTYPE_X1 value;
  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
  uint32_t typeKey;
  uint32_t x2Length;
};

extern "C" __global__ __aicore__ void addcmul(GM_ADDR input_data, GM_ADDR x1, GM_ADDR x2, GM_ADDR value, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  // TODO: user kernel impl
  KernelAddcmul op;

  op.Init(input_data, x1, x2, value, y, tiling_data.blockLength,
          tiling_data.tileNum, tiling_data.tileLength,
          tiling_data.lasttileLength, tiling_data.typeKey, tiling_data.x2Length);
  op.Process();
}

#ifndef __CCE_KT_TEST__
void addcmul_do(uint32_t blockDim, void* l2ctrl, void* stream,
                       uint8_t* input_data, uint8_t* x1, uint8_t* x2, uint8_t* value, uint8_t* y,
                       uint8_t* workspace, uint8_t* tiling) {
  addcmul<<<blockDim, l2ctrl, stream>>>(input_data, x1, x2, value, y, workspace, tiling);
}
#endif